"""词向量和数字向量的转换"""


class NumSeq:
    PAD_TAG = "PAD"
    UNK_TAG = "UNK"
    EOS_TAG = "EOS"
    SOS_TAG = "SOS"

    PAD = 0
    UNK = 1
    EOS = 2
    SOS = 3

    dict = {
        PAD_TAG: PAD,
        UNK_TAG: UNK,
        EOS_TAG: EOS,
        SOS_TAG: SOS
    }

    reverse_dict = {}

    def __init__(self):
        for i in range(10):
            self.dict[str(i)] = len(self.dict)

        self.reverse_dict = dict(zip(self.dict.values(), self.dict.keys()))

    def word_to_num_transform(self, sentence, max_len=None, add_eos=False):
        """
        将词向量转换为数字向量
        :param sentence: 按词分割好的句子
        :param max_len: 向量的最大长度
        :param add_eos: 添加开始标识
        :return:
        """
        if len(sentence) >= max_len:
            sentence = sentence[:max_len]
        sentence_len = len(sentence)
        if add_eos:
            sentence += [self.EOS_TAG]
        if sentence_len < max_len:
            sentence += [self.PAD_TAG] * (max_len - sentence_len)

        return [self.dict.get(word, self.UNK) for word in sentence]

    def num_to_word_transform(self, num_vector):
        """将数字向量转换为词向量"""
        return [self.reverse_dict.get(num, self.UNK) for num in num_vector]

    def __len__(self):
        return len(self.dict)


if __name__ == '__main__':
    num_seq = NumSeq()
    print(num_seq.word_to_num_transform(['1', '2', '4'], max_len=10))
    print(num_seq.num_to_word_transform([1, 2, 4], max_len=10))



